#!/usr/bin/env python
import os
import os.path
import re

from threading import Thread, Lock
from Queue import Queue, Empty

from Ganga.Utility.Shell import Shell
from GangaAtlas.Lib.AtlasLCGRequirements.AtlasLCGRequirements import getCEsForSites
from GangaAtlas.Lib.Athena.dm_util import get_transfer_protocols, get_srm_host

logger = Ganga.Utility.logging.getLogger(modulename=False)

def resolve_atlas_rtag(atlas_project='', atlas_release='', platform='i686-slc5-gcc43-opt'):

    revs = atlas_release.split('.')

    atlas_rtag        = ''
    atlas_rtag_suffix = ''

    if len(revs) > 3:
        if not atlas_project:
            atlas_project = 'AtlasProduction'

        if atlas_project == 'AtlasProduction':
            ## AtlasProduction takes only first 3 versioning numbers
            atlas_rtag_suffix = '%s-%s' % ('.'.join(revs[0:4]), platform)
        else:
            atlas_rtag_suffix = '%s-%s' % ('.'.join(revs), platform)
    else:
        ## versioning number not more than 3, assuming the offline version
        atlas_project = ''
        atlas_rtag_suffix = '%s-%s' % ('.'.join(revs[0:3]), platform)

    if atlas_project in [ 'AtlasPoint1', 'AtlasTier0', 'AtlasProduction' ]:
        atlas_rtag_prefix = 'VO-atlas-%s' % re.sub('^Atlas', '', atlas_project).lower()
    elif atlas_project:
        atlas_rtag_prefix = 'VO-atlas-%s' % atlas_project.lower()
    else:
        atlas_rtag_prefix = 'VO-atlas-offline'

    atlas_rtag = '%s-%s' % (atlas_rtag_prefix, atlas_rtag_suffix)

    logger.debug('Atlas release Tag: %s' % atlas_rtag)

    return atlas_rtag

def ama_mkstruct_jobs(amaConfig='', amaeventDsList=[], amaFlags={}, maxFilesPerJob=10, nthread=1, dsloc_pref_regex=''):
    """
    generating grid jobs to make struct ntuples given a list of AMAEvent dataset containers
    """

    import ddm_util
    import ama_util
    #import common_util

    jobList = []

    sampleList = {}

    if len(amaeventDsList) < nthread:
        nthread = len(amaeventDsList)
        
    wq = Queue(len(amaeventDsList))

    for ds in amaeventDsList:
        wq.put(ds)

    mylock = Lock()

    def worker(id):
        while not wq.empty():
            try:
                ds = wq.get(block=True, timeout=1)

                if ds.endswith('/'):
                    sampleName = ds[:-1]
                    mySampleList = ama_util.make_fs_sample(ds, dsloc_pref_regex=dsloc_pref_regex, sampleName=sampleName, chunkSize=maxFilesPerJob)

                    if mySampleList:

                        mylock.acquire()

                        for loc in mySampleList.keys():
                            if loc not in sampleList.keys():
                                sampleList[loc] = []
                            sampleList[loc] += mySampleList[loc]

                        mylock.release()

            except Empty:
                pass

    # generating sample definition files in parallel
    threads   = []
    for i in range(nthread):
        t = Thread(target=worker, kwargs={'id': i})
        t.setDaemon(False)
        threads.append(t)

    for t in threads:
        t.start()

    for t in threads:
        t.join()

    for loc in sampleList.keys():

        ## 1. resolve the close CE
        celist = getCEsForSites( [loc], excluded_ids=[], CREAM=False, cputime=0)

        ## 2. resolve the local data access protocol supported by the SE

        my_protocol = 'lcgcp'

        protocols = get_transfer_protocols( get_srm_host(loc) )

        try:
            protocols.remove('gsiftp')
        except ValueError:
            pass

        if protocols:
            if 'file' in protocols:
                my_protocol = 'file'
            else:
                my_protocol = protocols[0]

        ## 3. call ama_ganga_job to make the Ganga job
        j = ama_ganga_job(amaConfig=amaConfig, sampleList=sampleList[loc], maxEvt=-1, amaFlags=amaFlags, platform='grid', gridCEList=celist, fsProtocol=my_protocol)

        if j:
            jobList.append(j)

    return jobList


def ama_ganga_job(amaConfig='', sampleList=[], maxEvt=-1, amaFlags={}, platform='grid', gridCEList=[], fsProtocol='lcgcp'):

    import Ganga.Core.Sandbox as Sandbox

    ## definition of local helper routines
    def setup_grid_proxy():
        '''internal routine to copy local grid proxy to the home directory (presuming it is a shared fs)'''

        ui_proxy_path = gridProxy.location()

        shr_proxy_path = os.path.join( os.path.join( os.path.expanduser('~'), '.globus'), '.proxy' )

        cmd = 'cp %s %s' % ( ui_proxy_path, shr_proxy_path )

        s = Shell()

        rc,out,m = s.cmd1( cmd, allowed_exit=[0,255] )

        if rc == 0:
            return shr_proxy_path
        else:
            return None
    ## end of definition of local helper routines

    if not amaConfig:
        logger.error('invalid amaConfig: %s' % repr(amaConfig))
        return None

    if not sampleList:
        logger.error('invalid sampleList: %s' % repr(sampleList))
        return None

    if platform not in ['local','stb','grid']:
        logger.error('platform not supported: %s' % platform)
        return None

    for f in sampleList:
        if not os.path.exists(f):
            logger.error('file not found: %s' % f)
            return None

    ## resolve AtlasVersion and AtlasProject env. variables
    ##  - AtlasVersion is mandatory
    ##  - AtlasProject is optional (if not set, presuming AtlasOffline)
    AtlasVersion = ''
    try:
        AtlasVersion = os.environ['AtlasVersion']
    except:
        pass
    if not AtlasVersion:
        logger.error('AtlasVersion not defined')
        return None

    AtlasProject = ''
    try:
        AtlasProject = os.environ['AtlasProject']
    except:
        pass

    ## resolve TestArea and make a tarbll containing AMACore/User libraries 
    TestArea = ''
    try:
        TestArea = os.environ['TestArea']
    except:
        pass

    if (not TestArea) or (not os.path.exists(TestArea)):
        logger.error('TestArea not found')
        return None


    ama_util = '/project/atlas/nikhef/ganga/tools/amauser/ama_util.sh'
    ama_exe  = '/project/atlas/nikhef/ganga/tools/amauser/ama_exe.sh'
    fs_cmd   = '/project/atlas/nikhef/ganga/tools/amauser/lcg-cp'

    j = Job()

    ## pack the ama binary in a tarball
    ama_ia_dirs    = [os.environ['CMTCONFIG']]
    ama_ia_tgz     = os.path.join(j.inputdir, 'amauser.tgz')
    ama_ia_pkg_cmd = 'cd %s; tar czhf %s %s' % ( TestArea, ama_ia_tgz, ' '.join(map(lambda x:os.path.join('InstallArea',x), ama_ia_dirs)) )

    logger.debug(ama_ia_pkg_cmd)

    s = Shell()
    rc,out,m = s.cmd1( ama_ia_pkg_cmd, allowed_exit=[0,255] )

    if rc != 0:
        print out
        return None

    ## set application
    my_platform = platform
    if platform in ['local','stb']:
        my_platform = 'local'

    j.application      = Executable()
    j.application.exe  = File(ama_exe)
    j.application.args = []
    j.application.env  = {'ATLAS_PRODUCTION'     : AtlasVersion,
                          'ATLAS_PROJECT'        : AtlasProject,
                          'ATLAS_RELEASE'        : '.'.join( AtlasVersion.split('.')[0:3] ),
                          'FILE_STAGER_PROTOCOL' : fsProtocol,
                          'AMAUSER_PLATFORM'     : my_platform}

    if platform in ['local','stb']:
        j.application.env['GANGA_GLITE_UI'] = config['LCG']['GLITE_SETUP']
        j.application.env['ATLAS_SOFTWARE'] = config['Athena']['ATLAS_SOFTWARE']    

    ## set splitter
    j.splitter            = GenericSplitter()
    j.splitter.attribute  = 'application.args'
    j.splitter.values     =  []

    amaFlagList = []
    for k,v in amaFlags.items():
        amaFlagList.append( '%s=%s' % (k, repr(v)) )

    for f in sampleList:
        j.splitter.values.append( [ os.path.basename(f), os.path.basename(amaConfig), repr(maxEvt) ] + amaFlagList )
    
    ## create a tarball containing all sample definitions
    inws = j.__dict__['_impl'].getInputWorkspace()

    j.inputsandbox   = [File(ama_util)]
    j.inputsandbox  += [File(fs_cmd)]
    j.inputsandbox  += [File(amaConfig)]
    j.inputsandbox  += [File(ama_ia_tgz)]

    ## add all sample file definitions in input sandbox
    ## there is no way to attach individual one to a subjob accordingly
    for f in sampleList:
        j.inputsandbox += [File(f)]

    j.outputsandbox += ['summary/*', 'FileStager.*']

    ## set backend and force it to run at NIKHEF CE
    if platform == 'stb':
        ## set backend to PBS for submitting batch jobs to Stoomboot nodes
 
        shr_gproxy_path = setup_grid_proxy()

        if shr_gproxy_path:
            j.application.env['SHR_X509_USER_PROXY'] = shr_gproxy_path
            j.backend = PBS()
        else:
            logger.error('fail to copy local grid proxy to user HOME directory')
            j.remove()
            return None

    elif platform == 'local':
        ## set backend to local machine 
        j.backend = Local()
    elif platform == 'grid':
        ## set backend and force it to run at NIKHEF CE
        j.backend    = LCG()

        if not gridCEList:
            gridCEList = [ 'gazon.nikhef.nl:2119/jobmanager-pbs-medium' ]

        if len(gridCEList) == 1:
            j.backend.CE = gridCEList[0]
        else:
            j.backend.requirements = LCGRequirements()
            ce_requirement = ' ||\n     '.join([ 'other.GlueCEUniqueID == "%s"' % ce for ce in gridCEList ])
            j.backend.requirements.other.append( '(%s)' % ce_requirement )

            ## add atlas releast tag requirement
            atlas_rtag = resolve_atlas_rtag(atlas_project=AtlasProject, atlas_release=AtlasVersion, platform='i686-slc5-gcc43-opt')
            j.backend.requirements.software.append( atlas_rtag )
    else:
        logger.error('platform not supported: %s' % platform)
        j.remove()
        return None

    return j

## main program for test ##
#if __name__ == "__main__":
#
#    amaConfig  = 'Config/NewTtbar.conf'
#    sampleList = []
#    for i in range(0,6):
#        sampleList.append('samples/test_sample.%d.def' % i)
#    maxEvt = 1000
#
#    j = ama_ganga_job(amaConfig, sampleList, maxEvt)
